from PIL import Image
import numpy as np
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization

def to_binary(string):
    binary_list = []
    utf8_bytes = string.encode('utf-8')
    for byte in utf8_bytes:
        binary_list.append(bin(byte).replace('0b', '').zfill(8))
    return binary_list

# rsa修改
def bytes_to_binary(rsa_bytes):
    binary_list = []
    for byte in rsa_bytes:
        binary_list.append(bin(byte).replace('0b', '').zfill(8))
    return binary_list


def from_binary(binary_list):
    byte_array = bytearray(int(binary_list[i], 2) for i in range(len(binary_list)))
    result = byte_array.decode('utf-8')
    return result

# rsa修改
def binary_to_bytes(binary_list):
    # 将每个二进制字符串转换为整数
    byte_list = [int(binary_str, 2) for binary_str in binary_list]
    # 将整数列表转换为字节串
    rsa_bytes = bytes(byte_list)
    # 加载私钥  
    return rsa_bytes


class Write_LSB:  # LSB写数据类
    def __init__(self, pic_data):
        self.pic_data = np.array(pic_data)  # 图片数据
        

    def write_data(self, data_list):
        data = ''.join(data_list)
        data_len = len(data)
        for i in range(data_len):
            # 将数据嵌入到每个颜色通道的最低有效位
            pixel = self.pic_data[i // self.pic_data.shape[1], i % self.pic_data.shape[1]]
            if data[i] == '0':
                pixel[2] = pixel[2] & 0xfe
            else:
                pixel[2] = pixel[2] | 1
            self.pic_data[i // self.pic_data.shape[1], i % self.pic_data.shape[1]] = pixel

    def get_pic(self):  # 返回结果图片
        return self.pic_data

class Read_LSB:  # LSB读数据类
    def __init__(self, pic_data):
        self.pic_data = np.array(pic_data)  # 图片数据

    def read_data(self, data_len):
        data_list = []
        for i in range(data_len*8):
            # 从每个颜色通道的最低有效位提取数据
            pixel = self.pic_data[i // self.pic_data.shape[1], i % self.pic_data.shape[1]]
            data_value = str(pixel[2] & 1)  # 提取红、绿、蓝通道的最低有效位
            data_list.append(data_value)
        return ''.join(data_list)

def actor_lsb(path, data_r):
    data = Image.open(path)
    data = data.convert("RGB")
    data = np.array(data)
    a = Write_LSB(data)
    # rsa修改
    data_w = bytes_to_binary(data_r)
    a.write_data(data_w)
    r = a.get_pic()
    r = r.astype("uint8")
    r = Image.fromarray(r, "RGB")  # 修改为RGB模式
    r.save("ans_lsb.png")
    return len(data_w)  # 返回隐藏数据的大小

def reader_lsb(path, len_data):
    data = Image.open(path)
    data = data.convert("RGB")  # 修改为RGB模式
    data = np.array(data)
    b = Read_LSB(data)
    s = b.read_data(len_data)  # 提取比特数据
    return binary_to_bytes([s[i:i + 8] for i in range(0, len(s), 8)])